import numpy as np
import Image
import os
import h5py
import csv

def vec_imgs_from_txt_file(img_folder,csv_file,res_folder,batch_size=1000):
    #read csv file
    path_and_labels = []
    f = open(csv_file,'rb')
    for line in f:
        line = line.strip('\r\n')
        name, label = line.split(',')
        path = img_folder + '/' + name
        label = int(label)
        path_and_labels.append((path,label))
    f.close()

    #vec imgs
    img_arrs = []
    labels = []
    img_num = 0
    noise = 0
    img_convert = 0
    for path_and_label in path_and_labels:
        img_num += 1
        path, label = path_and_label
        img = Image.open(path)
        if img.mode == "CMYK":
            img = img.convert("RGB")
            img_convert += 1
        img_arr = np.asarray(img,dtype='string')
        img_arr = img_arr.transpose(2,0,1)
        if img_arr.shape != (3,64,64):
            noise += 1
            continue
        labels.append(label)
        img_arrs.append(img_arr)

        if img_num%1000 == 0:
            print 'noise: ' + str(noise),
            print '  img_num: ' + str(img_num),
            print '  img_convert: ' + str(img_convert) + '\r',
    print 'noise: ' + str(noise)
    print 'img_num: ' + str(img_num)
    print 'img_convert: ' + str(img_convert)
    img_arrs = np.asarray(img_arrs,dtype='float32')
    labels = np.asarray(labels,dtype='int32')


    if not res_folder.endswith('/'):
        res_folder += '/'
    if not os.path.exists(res_folder):
        os.mkdir(res_folder)
    x = img_arrs
    y= labels
    n_batch = len(x) / batch_size
    i = 0
    for i in range(n_batch):
        print 'vector_num: ' + str((i+1)*batch_size) + '\r',;
        file_name = res_folder + str(i) + '.hdf5'
        batch_x = x[ i*batch_size: (i+1)*batch_size]
        batch_y = y[ i*batch_size: (i+1)*batch_size]
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.create_dataset('y',data=batch_y)
        f.close()
    if n_batch * batch_size < len(x):
        batch_x = x[n_batch*batch_size: ]
        batch_y = y[n_batch*batch_size: ]
        file_name = res_folder + str(n_batch) + '.hdf5'
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.create_dataset('y',data=batch_y)
        f.close()
    print 'vector_num: ' + str((i+1)*batch_size)
def test_vec_imgs_from_folder(img_folder,res_folder,name_file,batch_size=1000):
    f = file(name_file,'wb')
    csvwriter = csv.writer(f)
    img_arrs = []
    img_num = 0
    noise = 0
    img_convert = 0
    for img in os.listdir(img_folder):
        csvwriter.writerow([img])

        img_num += 1
        path = img_folder + '/' + img
        img = Image.open(path)
        if img.mode == "CMYK":
            img = img.convert("RGB")
            img_convert += 1
        img_arr = np.asarray(img,dtype='string')
        img_arr = img_arr.transpose(2,0,1)
        if img_arr.shape != (3,64,64):
            noise += 1
            continue
        img_arrs.append(img_arr)
        if img_num%1000 == 0:
            print 'noise: ' + str(noise),
            print '  img_num: ' + str(img_num),
            print '  img_convert: ' + str(img_convert) + '\r',
    f.close()
    print 'noise: ' + str(noise)
    print 'img_num: ' + str(img_num)
    print 'img_convert: ' + str(img_convert)


    img_arrs = np.asarray(img_arrs,dtype='float32')


    if not res_folder.endswith('/'):
        res_folder += '/'
    if not os.path.exists(res_folder):
        os.mkdir(res_folder)
    x = img_arrs
    n_batch = len(x) / batch_size
    i = 0
    for i in range(n_batch):
        print 'vector_num: ' + str((i+1)*batch_size) + '\r',;
        file_name = res_folder + str(i) + '.hdf5'
        batch_x = x[ i*batch_size: (i+1)*batch_size]
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.close()
    if n_batch * batch_size < len(x):
        batch_x = x[n_batch*batch_size: ]
        file_name = res_folder + str(n_batch) + '.hdf5'
        f = h5py.File(file_name,'w')
        f.create_dataset('x',data=batch_x)
        f.close()
    print 'vector_num: ' + str((i+1)*batch_size)

print 'vec_imgs_from_csv_file()'
"""
print 'train'
vec_imgs_from_txt_file(img_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train64',
                       csv_file='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/train_label.txt',
                       res_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/ProData/64/submit/trainVec',
                       batch_size=10000)"""
print 'test'
test_vec_imgs_from_folder(img_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/test64',
                          res_folder='/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/JD/ProData/64/submit/testVec',
                          name_file='/home/dell/wxm/Code/JD/log_records/submit/baseline/64/test_names.csv',
                          batch_size=10000)
